"""
@author: linruihan
@file: SMPConverter.py
@time: 2024/11/25
@desc: 用于破解smp加密音乐文件，参考https://www.52pojie.cn/thread-859330-1-1.html网址。
"""
import argparse
import os


def xor_circular(data: bytes, key: int) -> bytes:
    """
    对数据进行异或运算
    :param data: 文件数据
    :param key: 密钥
    :return: 异或后的数据
    """
    k_len = int(len(hex(key)) / 2 - 1)
    r = bytearray()
    for i in range(len(data)):
        # 逐字节异或
        r.append(data[i] ^ key.to_bytes(4, byteorder='big')[i % k_len])
    return bytes(r)


def smp_decoder(file_path: str, key: int, output_path: str):
    """
    解密smp文件
    :param file_path: smp文件路径
    :param key: 密钥
    :param output_path: 输出文件路径
    """
    with open(file_path, 'rb') as fi:
        data = fi.read()
        result = xor_circular(data, key)

        if output_path.lower().endswith('.mp3') is True and test_mp3(result) is False:
            raise ValueError('decode invalid mp3 file')

        if '/' in output_path:
            os.makedirs(os.path.dirname(output_path), exist_ok=True)

        with open(output_path, 'wb') as fo:
            fo.write(result)


def smp_encoder(file_path: str, key: int, output_path: str):
    """
    加密smp文件
    :param file_path: mp3文件路径
    :param key: 密钥
    :param output_path: 输出文件路径
    """
    with open(file_path, 'rb') as fi:
        data = fi.read()
        result = xor_circular(data, key)

        if '/' in output_path:
            os.makedirs(os.path.dirname(output_path), exist_ok=True)

        with open(output_path, 'wb') as fo:
            fo.write(result)


def test_mp3(data):
    """
    测试mp3文件正确性
    :param data: mp3文件数据
    """
    return data[:3] == 'ID3'.encode()


def main():
    parser = argparse.ArgumentParser(description='SMP Converter')
    parser.add_argument('-m', '--mode', choices=['decode', 'encode'], default='decode', help='mode')
    parser.add_argument('-i', '--input', help='input file', required=True)
    parser.add_argument('-k', '--key', help='key (hexadecimal string)', required=True)
    parser.add_argument('-o', '--output', help='output file', required=True)
    args = parser.parse_args()

    try:
        k = int(args.key, 16)
    except ValueError:
        print('Error: Invalid key')
        exit(1)

    if os.path.exists(args.input) is False:
        print('Error: Input File not found')
        exit(1)

    if args.mode == 'decode':
        try:
            smp_decoder(args.input, k, args.output)
        except ValueError:
            print('Error: Input smp file with key decode invalid mp3 file, please check input file and key!')
            exit(1)
    else:
        smp_encoder(args.input, k, args.output)


if __name__ == '__main__':
    main()
